#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import os
import cv2
import numpy as np
import time
from PyQt5.QtWidgets import *
from PyQt5.QtGui import QPixmap, QImage, QFont
from PyQt5.QtCore import QTimer, pyqtSignal, QThread, Qt

# 添加上级目录到Python路径以导入MvImport模块
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

from MvImport.MvCameraControl_class import *
from MvImport.MvErrorDefine_const import *
from MvImport.PixelType_header import *

# 尝试导入YOLO和OCR
try:
    from ultralytics import YOLO
    YOLO_AVAILABLE = True
except ImportError:
    YOLO_AVAILABLE = False
    print("Warning: YOLO not available. Install with: pip install ultralytics")

try:
    from paddleocr import PaddleOCR
    PADDLEOCR_AVAILABLE = True
except ImportError:
    PADDLEOCR_AVAILABLE = False
    print("Warning: PaddleOCR not available. Install with: pip install paddleocr")

class RecognitionThread(QThread):
    """识别线程"""
    result_ready = pyqtSignal(str, float, np.ndarray)  # 识别结果, 置信度, 标注图像
    
    def __init__(self, camera_index=0):
        super().__init__()
        self.camera_index = camera_index
        self.is_running = False
        self.cam = None
        
        # 初始化识别模型
        self.detector = None
        self.ocr = None
        self.init_models()
        
    def init_models(self):
        """初始化识别模型"""
        # 初始化YOLO检测器
        if YOLO_AVAILABLE:
            try:
                # 尝试加载训练好的模型，如果没有则使用预训练模型
                if os.path.exists('best.pt'):
                    self.detector = YOLO('best.pt')
                    print("加载自定义YOLO模型: best.pt")
                else:
                    self.detector = YOLO('yolov8n.pt')
                    print("加载预训练YOLO模型: yolov8n.pt")
            except Exception as e:
                print(f"YOLO模型加载失败: {e}")
                self.detector = None
        
        # 初始化PaddleOCR
        if PADDLEOCR_AVAILABLE:
            try:
                self.ocr = PaddleOCR(use_angle_cls=True, lang='en', show_log=False)
                print("PaddleOCR初始化成功")
            except Exception as e:
                print(f"PaddleOCR初始化失败: {e}")
                self.ocr = None
    
    def run(self):
        """运行识别"""
        try:
            # 枚举设备
            deviceList = MV_CC_DEVICE_INFO_LIST()
            tlayerType = MV_GIGE_DEVICE | MV_USB_DEVICE
            
            ret = MvCamera.MV_CC_EnumDevices(tlayerType, deviceList)
            if ret != 0 or deviceList.nDeviceNum == 0:
                print("没有找到设备!")
                return
                
            if self.camera_index >= deviceList.nDeviceNum:
                print(f"设备索引 {self.camera_index} 超出范围!")
                return
            
            # 创建相机实例
            self.cam = MvCamera()
            
            # 选择设备并创建句柄
            stDeviceList = cast(deviceList.pDeviceInfo[self.camera_index], POINTER(MV_CC_DEVICE_INFO)).contents
            
            ret = self.cam.MV_CC_CreateHandle(stDeviceList)
            if ret != 0:
                print(f"创建句柄失败! ret[0x{ret:08x}]")
                return
                
            # 打开设备
            ret = self.cam.MV_CC_OpenDevice(MV_ACCESS_Exclusive, 0)
            if ret != 0:
                print(f"打开设备失败! ret[0x{ret:08x}]")
                self.cam.MV_CC_DestroyHandle()
                return
            
            # 设置触发模式为off
            ret = self.cam.MV_CC_SetEnumValue("TriggerMode", MV_TRIGGER_MODE_OFF)
            if ret != 0:
                print(f"设置触发模式失败! ret[0x{ret:08x}]")
            
            # 开始取流
            ret = self.cam.MV_CC_StartGrabbing()
            if ret != 0:
                print(f"开始取流失败! ret[0x{ret:08x}]")
                self.cam.MV_CC_CloseDevice()
                self.cam.MV_CC_DestroyHandle()
                return
            
            print("相机开始采集...")
            self.is_running = True
            
            # 获取图像数据
            stOutFrame = MV_FRAME_OUT()
            memset(byref(stOutFrame), 0, sizeof(stOutFrame))
            
            last_recognition_time = 0
            recognition_interval = 1.0  # 每1秒识别一次
            
            while self.is_running:
                ret = self.cam.MV_CC_GetImageBuffer(stOutFrame, 1000)
                if ret == 0:
                    # 转换图像数据
                    data = (c_ubyte * stOutFrame.stFrameInfo.nFrameLen)()
                    cdll.msvcrt.memcpy(byref(data), stOutFrame.pBufAddr, stOutFrame.stFrameInfo.nFrameLen)
                    
                    # 根据像素格式转换图像
                    if stOutFrame.stFrameInfo.enPixelType == PixelType_Gvsp_Mono8:
                        # 单色图像
                        img_data = np.frombuffer(data, dtype=np.uint8)
                        img = img_data.reshape((stOutFrame.stFrameInfo.nHeight, stOutFrame.stFrameInfo.nWidth))
                        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
                    elif stOutFrame.stFrameInfo.enPixelType == PixelType_Gvsp_RGB8_Packed:
                        # RGB图像
                        img_data = np.frombuffer(data, dtype=np.uint8)
                        img = img_data.reshape((stOutFrame.stFrameInfo.nHeight, stOutFrame.stFrameInfo.nWidth, 3))
                        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
                    else:
                        # 其他格式，尝试转换为BGR
                        img_data = np.frombuffer(data, dtype=np.uint8)
                        img = img_data.reshape((stOutFrame.stFrameInfo.nHeight, stOutFrame.stFrameInfo.nWidth, -1))
                        if len(img.shape) == 2:
                            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
                    
                    # 控制识别频率
                    current_time = time.time()
                    if current_time - last_recognition_time > recognition_interval:
                        # 执行识别
                        result, confidence, annotated_img = self.perform_recognition(img.copy())
                        if result:
                            self.result_ready.emit(result, confidence, annotated_img)
                        last_recognition_time = current_time
                    
                    # 释放图像缓存
                    self.cam.MV_CC_FreeImageBuffer(stOutFrame)
                else:
                    time.sleep(0.01)
                    
        except Exception as e:
            print(f"识别线程错误: {e}")
        finally:
            self.stop()
    
    def perform_recognition(self, image):
        """执行识别"""
        try:
            result_text = ""
            confidence = 0.0
            annotated_img = image.copy()
            
            # 1. YOLO目标检测
            if self.detector:
                results = self.detector(image, verbose=False)
                
                for result in results:
                    boxes = result.boxes
                    if boxes is not None:
                        for box in boxes:
                            # 获取边界框
                            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                            conf = box.conf[0].cpu().numpy()
                            
                            # 绘制检测框
                            cv2.rectangle(annotated_img, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
                            cv2.putText(annotated_img, f"Object: {conf:.2f}", (int(x1), int(y1)-10), 
                                      cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                            
                            # 提取ROI区域
                            roi = image[int(y1):int(y2), int(x1):int(x2)]
                            
                            # 2. OCR识别
                            if self.ocr and roi.size > 0:
                                ocr_result = self.ocr.ocr(roi, cls=True)
                                
                                if ocr_result and ocr_result[0]:
                                    for line in ocr_result[0]:
                                        text = line[1][0]
                                        ocr_conf = line[1][1]
                                        
                                        # 过滤数字
                                        if text.isdigit() and len(text) <= 3:  # 假设数字不超过3位
                                            result_text = text
                                            confidence = ocr_conf
                                            
                                            # 在ROI上标注识别结果
                                            cv2.putText(annotated_img, f"Number: {text}", 
                                                      (int(x1), int(y2)+20), cv2.FONT_HERSHEY_SIMPLEX, 
                                                      0.7, (0, 0, 255), 2)
                                            break
            
            # 如果没有检测到物体，尝试直接OCR
            if not result_text and self.ocr:
                ocr_result = self.ocr.ocr(image, cls=True)
                if ocr_result and ocr_result[0]:
                    for line in ocr_result[0]:
                        text = line[1][0]
                        ocr_conf = line[1][1]
                        
                        if text.isdigit() and len(text) <= 3:
                            result_text = text
                            confidence = ocr_conf
                            break
            
            return result_text, confidence, annotated_img
            
        except Exception as e:
            print(f"识别过程错误: {e}")
            return "", 0.0, image
    
    def stop(self):
        """停止识别"""
        self.is_running = False
        if self.cam:
            try:
                self.cam.MV_CC_StopGrabbing()
                self.cam.MV_CC_CloseDevice()
                self.cam.MV_CC_DestroyHandle()
            except:
                pass

class ObjectRecognitionDemo(QMainWindow):
    """物体识别演示窗口"""
    
    def __init__(self):
        super().__init__()
        self.recognition_thread = None
        self.init_ui()
        
    def init_ui(self):
        """初始化界面"""
        self.setWindowTitle('海康威视相机 - 物体识别系统')
        self.setGeometry(100, 100, 1200, 800)
        
        # 创建中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 创建主布局
        main_layout = QHBoxLayout(central_widget)
        
        # 左侧：图像显示区域
        left_widget = QWidget()
        left_layout = QVBoxLayout(left_widget)
        
        # 图像显示标签
        self.image_label = QLabel()
        self.image_label.setMinimumSize(800, 600)
        self.image_label.setStyleSheet("border: 2px solid gray;")
        self.image_label.setText("等待连接相机...")
        self.image_label.setAlignment(Qt.AlignCenter)
        left_layout.addWidget(self.image_label)
        
        main_layout.addWidget(left_widget)
        
        # 右侧：控制面板
        right_widget = QWidget()
        right_widget.setMaximumWidth(300)
        right_layout = QVBoxLayout(right_widget)
        
        # 标题
        title_label = QLabel("识别控制面板")
        title_label.setFont(QFont("Arial", 14, QFont.Bold))
        title_label.setAlignment(Qt.AlignCenter)
        right_layout.addWidget(title_label)
        
        # 控制按钮
        self.start_button = QPushButton("开始识别")
        self.start_button.setMinimumHeight(40)
        self.start_button.clicked.connect(self.start_recognition)
        right_layout.addWidget(self.start_button)
        
        self.stop_button = QPushButton("停止识别")
        self.stop_button.setMinimumHeight(40)
        self.stop_button.clicked.connect(self.stop_recognition)
        self.stop_button.setEnabled(False)
        right_layout.addWidget(self.stop_button)
        
        # 识别结果显示
        result_group = QGroupBox("识别结果")
        result_layout = QVBoxLayout(result_group)
        
        self.result_label = QLabel("等待识别...")
        self.result_label.setFont(QFont("Arial", 12))
        self.result_label.setAlignment(Qt.AlignCenter)
        self.result_label.setStyleSheet("background-color: #f0f0f0; padding: 10px; border: 1px solid #ccc;")
        result_layout.addWidget(self.result_label)
        
        self.confidence_label = QLabel("置信度: --")
        self.confidence_label.setFont(QFont("Arial", 10))
        result_layout.addWidget(self.confidence_label)
        
        right_layout.addWidget(result_group)
        
        # 状态信息
        status_group = QGroupBox("系统状态")
        status_layout = QVBoxLayout(status_group)
        
        self.status_label = QLabel("就绪")
        status_layout.addWidget(self.status_label)
        
        # 模型状态
        yolo_status = "✅ YOLO可用" if YOLO_AVAILABLE else "❌ YOLO不可用"
        ocr_status = "✅ PaddleOCR可用" if PADDLEOCR_AVAILABLE else "❌ PaddleOCR不可用"
        
        self.yolo_status_label = QLabel(f"YOLO: {yolo_status}")
        self.ocr_status_label = QLabel(f"PaddleOCR: {ocr_status}")
        
        status_layout.addWidget(self.yolo_status_label)
        status_layout.addWidget(self.ocr_status_label)
        
        right_layout.addWidget(status_group)
        
        # 添加弹性空间
        right_layout.addStretch()
        
        main_layout.addWidget(right_widget)
        
        self.current_image = None
        
    def start_recognition(self):
        """开始识别"""
        if self.recognition_thread is None or not self.recognition_thread.is_running:
            self.recognition_thread = RecognitionThread(0)
            self.recognition_thread.result_ready.connect(self.update_result)
            self.recognition_thread.start()
            
            self.start_button.setEnabled(False)
            self.stop_button.setEnabled(True)
            self.status_label.setText("正在识别...")
    
    def stop_recognition(self):
        """停止识别"""
        if self.recognition_thread and self.recognition_thread.is_running:
            self.recognition_thread.stop()
            self.recognition_thread.wait()
            
            self.start_button.setEnabled(True)
            self.stop_button.setEnabled(False)
            self.status_label.setText("已停止")
    
    def update_result(self, result, confidence, annotated_img):
        """更新识别结果"""
        if result:
            self.result_label.setText(f"识别数字: {result}")
            self.confidence_label.setText(f"置信度: {confidence:.2f}")
            
            # 更新图像显示
            self.update_image(annotated_img)
        else:
            self.result_label.setText("未识别到数字")
            self.confidence_label.setText("置信度: --")
    
    def update_image(self, img):
        """更新图像显示"""
        if img is None:
            return
            
        # 确保图像是3通道的
        if len(img.shape) == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        
        # 调整图像大小以适应显示区域
        height, width, channel = img.shape
        max_width = self.image_label.width() - 10
        max_height = self.image_label.height() - 10
        
        if width > max_width or height > max_height:
            scale = min(max_width / width, max_height / height)
            new_width = int(width * scale)
            new_height = int(height * scale)
            img = cv2.resize(img, (new_width, new_height))
            height, width, channel = img.shape
        
        # 转换为QImage
        bytes_per_line = 3 * width
        q_img = QImage(img.data, width, height, bytes_per_line, QImage.Format_RGB888)
        
        # 显示图像
        pixmap = QPixmap.fromImage(q_img)
        self.image_label.setPixmap(pixmap)
    
    def closeEvent(self, event):
        """关闭事件"""
        self.stop_recognition()
        event.accept()

def main():
    app = QApplication(sys.argv)
    window = ObjectRecognitionDemo()
    window.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main() 